import cv2
import numpy as np
from ultralytics import YOLO
from collections import defaultdict

# 加载模型
model = YOLO("yolov8s-pose.pt")  # 加载官方模型
model = YOLO("./pt/v8s-best.pt")  # 加载自定义模型

# 配置视频相关的
video_path = "../materials/dog25fps.mp4"
# video_path = "../petvideo/dog25fps.mp4"
result_path = "../detectvideo/track_dog-v8s.mp4"

# 记录所有的id点信息
track_history = defaultdict(lambda: [])

if __name__ == '__main__':
    # 抓取视频
    capture = cv2.VideoCapture(video_path)
    if not capture.isOpened():
        print("Error opening Video file.")
        exit()

    # 视频参数
    fps = capture.get(cv2.CAP_PROP_FPS)
    frame_width = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
    frame_height = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)

    # 窗口等待时间
    wtime = (int)((1 / fps) * 1000)

    videoWriter = None

    while True:
        successCode, frame = capture.read()
        if not successCode:
            print("视频读取结束")
            break

        # 预测单张图片
        results = model.track(frame, persist=True)
        a_frame = results[0].plot()
        a_frame = np.copy(a_frame)

        # 所有id的序列号信息
        boxes = results[0].boxes.xywh.cpu()
        if results[0].boxes.id is None:
            track_ids = []
        else:
            track_ids = results[0].boxes.id.int().cpu().tolist()

        # 迭代每一个预测到的对象
        for box, track_id in zip(boxes, track_ids):

            # 获得框的大小
            x, y, w, h = box
            track = track_history[track_id]
            # 添加坐标
            track.append((float(x), float(y)))
            if len(track) > 75:
                track.pop(0)

            # 创建点实例
            points = np.hstack(track).astype(np.int32).reshape(-1, 1, 2)
            # 画到预测结果上
            cv2.polylines(a_frame, [points], isClosed=False, color=(0, 0, 255), thickness=3)

        # 初始化并写入视频文件
        if videoWriter is None:
            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
            videoWriter = cv2.VideoWriter(result_path, fourcc, fps, (int(frame_width), int(frame_height)))

        videoWriter.write(a_frame)

    # 释放资源
    capture.release()
    videoWriter.release()
    cv2.destroyAllWindows()
